#!/usr/bin/env python3
"""
MyLangExtract uv 集成示例
演示如何在使用 uv 管理的项目中集成 MyLangExtract
"""
import os
import requests
import json
from typing import Dict, List, Optional

class MyLangExtractClient:
    """
    MyLangExtract 客户端包装器
    适用于 uv 管理的项目集成
    """
    
    def __init__(self, api_key: str = None, provider: str = "zhipu"):
        """
        初始化客户端
        
        Args:
            api_key: API 密钥，如果不提供则从环境变量获取
            provider: 大模型提供商，默认为 zhipu
        """
        self.api_key = api_key or os.environ.get('ZHIPU_API_KEY')
        self.provider = provider
        
        # 提供商配置
        self.provider_configs = {
            "zhipu": {
                "base_url": "https://open.bigmodel.cn/api/paas/v4/chat/completions",
                "model": "glm-4"
            },
            "moonshot": {
                "base_url": "https://api.moonshot.cn/v1/chat/completions",
                "model": "moonshot-v1-8k"
            },
            "openai": {
                "base_url": "https://api.openai.com/v1/chat/completions",
                "model": "gpt-3.5-turbo"
            }
        }
        
        if not self.api_key:
            raise ValueError(f"API 密钥未配置，请设置 {provider.upper()}_API_KEY 环境变量")
        
        if provider not in self.provider_configs:
            raise ValueError(f"不支持的提供商: {provider}")
    
    def extract_info(self, text: str, schema: Dict, prompt: str = None) -> Dict:
        """
        执行信息提取
        
        Args:
            text: 要提取信息的文本
            schema: 提取字段的 JSON Schema
            prompt: 自定义提示词，如果不提供则使用默认提示
            
        Returns:
            包含提取结果的字典
        """
        if not prompt:
            prompt = f"请从以下文本中提取信息：{text}"
        
        config = self.provider_configs[self.provider]
        
        headers = {
            'Authorization': f'Bearer {self.api_key}',
            'Content-Type': 'application/json'
        }
        
        data = {
            'model': config['model'],
            'messages': [{'role': 'user', 'content': prompt}],
            'tools': [{'type': 'function', 'function': schema}],
            'tool_choice': {'type': 'function', 'function': {'name': schema['name']}},
            'temperature': 0.1
        }
        
        try:
            response = requests.post(config['base_url'], headers=headers, json=data, timeout=30)
            
            if response.status_code == 200:
                result = response.json()
                
                if ('choices' in result and len(result['choices']) > 0 and
                    'message' in result['choices'][0] and 
                    'tool_calls' in result['choices'][0]['message']):
                    
                    tool_calls = result['choices'][0]['message']['tool_calls']
                    if tool_calls and len(tool_calls) > 0:
                        arguments = json.loads(tool_calls[0]['function']['arguments'])
                        return {
                            'success': True, 
                            'data': arguments,
                            'provider': self.provider,
                            'model': config['model']
                        }
                
                return {'success': False, 'error': '未找到提取结果'}
            else:
                return {'success': False, 'error': f'API 错误: {response.status_code}'}
                
        except Exception as e:
            return {'success': False, 'error': str(e)}
    
    def extract_person_info(self, text: str) -> Dict:
        """提取人物信息的便捷方法"""
        schema = {
            "name": "extract_person",
            "description": "提取人物信息",
            "parameters": {
                "type": "object",
                "properties": {
                    "name": {"type": "string", "description": "姓名"},
                    "job": {"type": "string", "description": "职业"},
                    "age": {"type": "string", "description": "年龄"},
                    "location": {"type": "string", "description": "地点"},
                    "company": {"type": "string", "description": "公司"}
                },
                "required": ["name"]
            }
        }
        
        return self.extract_info(text, schema)
    
    def extract_company_info(self, text: str) -> Dict:
        """提取公司信息的便捷方法"""
        schema = {
            "name": "extract_company",
            "description": "提取公司信息",
            "parameters": {
                "type": "object",
                "properties": {
                    "company": {"type": "string", "description": "公司名称"},
                    "industry": {"type": "string", "description": "行业"},
                    "founded": {"type": "string", "description": "成立时间"},
                    "location": {"type": "string", "description": "总部地点"},
                    "ceo": {"type": "string", "description": "CEO"},
                    "employees": {"type": "string", "description": "员工数量"}
                },
                "required": ["company"]
            }
        }
        
        return self.extract_info(text, schema)
    
    def extract_product_info(self, text: str) -> Dict:
        """提取产品信息的便捷方法"""
        schema = {
            "name": "extract_product",
            "description": "提取产品信息",
            "parameters": {
                "type": "object",
                "properties": {
                    "product": {"type": "string", "description": "产品名称"},
                    "price": {"type": "string", "description": "价格"},
                    "features": {"type": "string", "description": "主要特性"},
                    "brand": {"type": "string", "description": "品牌"},
                    "category": {"type": "string", "description": "产品类别"}
                },
                "required": ["product"]
            }
        }
        
        return self.extract_info(text, schema)
    
    def batch_extract(self, texts: List[str], extract_type: str = "person") -> List[Dict]:
        """
        批量提取信息
        
        Args:
            texts: 文本列表
            extract_type: 提取类型 ("person", "company", "product")
            
        Returns:
            提取结果列表
        """
        extract_methods = {
            "person": self.extract_person_info,
            "company": self.extract_company_info,
            "product": self.extract_product_info
        }
        
        if extract_type not in extract_methods:
            raise ValueError(f"不支持的提取类型: {extract_type}")
        
        method = extract_methods[extract_type]
        results = []
        
        for text in texts:
            result = method(text)
            results.append(result)
        
        return results

def demo_usage():
    """演示使用方法"""
    print("MyLangExtract uv 集成示例")
    print("=" * 50)
    
    # 检查 API 密钥
    if not os.environ.get('ZHIPU_API_KEY'):
        print("⚠️  请先设置 ZHIPU_API_KEY 环境变量")
        print("export ZHIPU_API_KEY='your_api_key_here'")
        return
    
    try:
        # 初始化客户端
        client = MyLangExtractClient()
        
        print("✅ 客户端初始化成功")
        print(f"使用提供商: {client.provider}")
        
        # 示例 1: 提取人物信息
        print("\n🧪 示例 1: 人物信息提取")
        person_text = "张三是一名软件工程师，今年28岁，在北京的腾讯公司工作"
        result = client.extract_person_info(person_text)
        
        print(f"输入: {person_text}")
        if result['success']:
            print("✅ 提取成功:")
            for key, value in result['data'].items():
                if value:
                    print(f"  {key}: {value}")
        else:
            print(f"❌ 提取失败: {result['error']}")
        
        # 示例 2: 提取公司信息
        print("\n🏢 示例 2: 公司信息提取")
        company_text = "苹果公司成立于1976年，总部位于美国加利福尼亚州库比蒂诺，现任CEO是蒂姆·库克，员工超过15万人"
        result = client.extract_company_info(company_text)
        
        print(f"输入: {company_text}")
        if result['success']:
            print("✅ 提取成功:")
            for key, value in result['data'].items():
                if value:
                    print(f"  {key}: {value}")
        else:
            print(f"❌ 提取失败: {result['error']}")
        
        # 示例 3: 批量处理
        print("\n📦 示例 3: 批量处理")
        texts = [
            "李四是产品经理，30岁，在上海工作",
            "王五是设计师，25岁，在深圳的字节跳动公司",
            "赵六是数据分析师，27岁"
        ]
        
        results = client.batch_extract(texts, "person")
        
        for i, (text, result) in enumerate(zip(texts, results), 1):
            print(f"\n文本 {i}: {text}")
            if result['success']:
                print("  ✅ 提取成功:")
                for key, value in result['data'].items():
                    if value:
                        print(f"    {key}: {value}")
            else:
                print(f"  ❌ 提取失败: {result['error']}")
        
        print("\n" + "=" * 50)
        print("🎉 演示完成！")
        
    except Exception as e:
        print(f"❌ 演示失败: {e}")

def main():
    """主函数"""
    demo_usage()

if __name__ == "__main__":
    main()
